/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.gobblin.data.management.source;

import java.io.IOException;
import java.util.Iterator;
import java.util.List;
import java.util.Optional;
import java.util.Spliterator;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

import org.apache.gobblin.configuration.SourceState;
import org.apache.gobblin.configuration.WorkUnitState;
import org.apache.gobblin.dataset.Dataset;
import org.apache.gobblin.dataset.IterableDatasetFinder;
import org.apache.gobblin.dataset.PartitionableDataset;
import org.apache.gobblin.dataset.URNIdentified;
import org.apache.gobblin.dataset.comparators.URNLexicographicalComparator;
import org.apache.gobblin.runtime.task.NoopTask;
import org.apache.gobblin.source.workunit.BasicWorkUnitStream;
import org.apache.gobblin.source.workunit.WorkUnit;
import org.apache.gobblin.source.workunit.WorkUnitStream;

import com.google.common.collect.AbstractIterator;
import com.google.common.collect.Iterators;
import com.google.common.collect.Lists;
import com.google.common.collect.PeekingIterator;

import javax.annotation.Nullable;
import lombok.extern.slf4j.Slf4j;


/**
 * A source that processes datasets generated by a {@link org.apache.gobblin.dataset.DatasetsFinder}, processing a few of
 * them each run, and continuing from where it left off in the next run. When it is done processing all the datasets, it
 * starts over from the beginning. The datasets are processed in lexicographical order based on URN.
 *
 * TODO: handle retries
 */
@Slf4j
public abstract class LoopingDatasetFinderSource<S, D> extends DatasetFinderSource<S, D> {

  public static final String MAX_WORK_UNITS_PER_RUN_KEY = "gobblin.source.loopingDatasetFinderSource.maxWorkUnitsPerRun";
  public static final int MAX_WORK_UNITS_PER_RUN = 10;

  private static final String DATASET_URN = "gobblin.source.loopingDatasetFinderSource.datasetUrn";
  private static final String PARTITION_URN = "gobblin.source.loopingDatasetFinderSource.partitionUrn";
  private static final String WORK_UNIT_ORDINAL = "gobblin.source.loopingDatasetFinderSource.workUnitOrdinal";
  protected static final String END_OF_DATASETS_KEY = "gobblin.source.loopingDatasetFinderSource.endOfDatasets";

  private final URNLexicographicalComparator lexicographicalComparator = new URNLexicographicalComparator();

  /**
   * @param drilldownIntoPartitions if set to true, will process each partition of a {@link PartitionableDataset} as a
   *                                separate work unit.
   */
  public LoopingDatasetFinderSource(boolean drilldownIntoPartitions) {
    super(drilldownIntoPartitions);
  }

  @Override
  public List<WorkUnit> getWorkunits(SourceState state) {
    return Lists.newArrayList(getWorkunitStream(state).getMaterializedWorkUnitCollection());
  }

  @Override
  public WorkUnitStream getWorkunitStream(SourceState state) {
    try {
      int maxWorkUnits = state.getPropAsInt(MAX_WORK_UNITS_PER_RUN_KEY, MAX_WORK_UNITS_PER_RUN);

      List<WorkUnitState> previousWorkUnitStates = state.getPreviousWorkUnitStates();
      Optional<WorkUnitState> maxWorkUnit;
      try {
        maxWorkUnit = previousWorkUnitStates.stream().reduce((wu1, wu2) -> {
          int wu1Ordinal = wu1.getPropAsInt(WORK_UNIT_ORDINAL);
          int wu2Ordinal = wu2.getPropAsInt(WORK_UNIT_ORDINAL);
          return wu1Ordinal > wu2Ordinal ? wu1 : wu2;
        });
      } catch (NumberFormatException nfe) {
        throw new RuntimeException("Work units in state store are corrupted! Missing or malformed " + WORK_UNIT_ORDINAL);
      }

      String previousDatasetUrnWatermark = null;
      String previousPartitionUrnWatermark = null;
      if (maxWorkUnit.isPresent() && !maxWorkUnit.get().getPropAsBoolean(END_OF_DATASETS_KEY, false)) {
        previousDatasetUrnWatermark = maxWorkUnit.get().getProp(DATASET_URN);
        previousPartitionUrnWatermark = maxWorkUnit.get().getProp(PARTITION_URN);
      }

      IterableDatasetFinder datasetsFinder = createDatasetsFinder(state);

      Stream<Dataset> datasetStream = datasetsFinder.getDatasetsStream(Spliterator.SORTED, this.lexicographicalComparator);
      datasetStream = sortStreamLexicographically(datasetStream);

      return new BasicWorkUnitStream.Builder(new DeepIterator(datasetStream.iterator(), previousDatasetUrnWatermark,
          previousPartitionUrnWatermark, maxWorkUnits)).setFiniteStream(true).build();
    } catch (IOException ioe) {
      throw new RuntimeException(ioe);
    }
  }

  /**
   * A deep iterator that advances input streams until the correct position, then possibly iterates over partitions
   * of {@link PartitionableDataset}s.
   */
  private class DeepIterator extends AbstractIterator<WorkUnit> {
    private final Iterator<Dataset> baseIterator;
    private final int maxWorkUnits;

    private Iterator<PartitionableDataset.DatasetPartition> currentPartitionIterator;
    private int generatedWorkUnits = 0;

    public DeepIterator(Iterator<Dataset> baseIterator, String previousDatasetUrnWatermark,
        String previousPartitionUrnWatermark, int maxWorkUnits) {
      this.maxWorkUnits = maxWorkUnits;
      this.baseIterator = baseIterator;

      Dataset equalDataset = advanceUntilLargerThan(Iterators.peekingIterator(this.baseIterator), previousDatasetUrnWatermark);

      if (drilldownIntoPartitions && equalDataset != null && equalDataset instanceof PartitionableDataset) {
        this.currentPartitionIterator = getPartitionIterator((PartitionableDataset) equalDataset);
        advanceUntilLargerThan(Iterators.peekingIterator(this.currentPartitionIterator), previousPartitionUrnWatermark);
      } else {
        this.currentPartitionIterator = Iterators.emptyIterator();
      }
    }

    /**
     * Advance an iterator until the next value is larger than the reference.
     * @return the last value polled if it is equal to reference, or null otherwise.
     */
    @Nullable  private <T extends URNIdentified> T advanceUntilLargerThan(PeekingIterator<T> it, String reference) {
      if (reference == null) {
        return null;
      }

      int comparisonResult = -1;
      while (it.hasNext() && (comparisonResult = lexicographicalComparator.compare(it.peek(), reference)) < 0) {
        it.next();
      }
      return comparisonResult == 0 ? it.next() : null;
    }

    private Iterator<PartitionableDataset.DatasetPartition> getPartitionIterator(PartitionableDataset dataset) {
      try {
        return this.currentPartitionIterator = sortStreamLexicographically(
            dataset.getPartitions(Spliterator.SORTED, LoopingDatasetFinderSource.this.lexicographicalComparator)).iterator();
      } catch (IOException ioe) {
        log.error("Failed to get partitions for dataset " + dataset.getUrn());
        return Iterators.emptyIterator();
      }
    }

    @Override
    protected WorkUnit computeNext() {
      if (this.generatedWorkUnits >= this.maxWorkUnits) {
        return endOfData();
      }

      while (this.baseIterator.hasNext() || this.currentPartitionIterator.hasNext()) {
        if (this.currentPartitionIterator != null && this.currentPartitionIterator.hasNext()) {
          PartitionableDataset.DatasetPartition partition = this.currentPartitionIterator.next();
          WorkUnit workUnit = workUnitForDatasetPartition(partition);
          addDatasetInfoToWorkUnit(workUnit, partition.getDataset(), this.generatedWorkUnits++);
          addPartitionInfoToWorkUnit(workUnit, partition);
          return workUnit;
        }

        Dataset dataset = this.baseIterator.next();
        if (drilldownIntoPartitions && dataset instanceof PartitionableDataset) {
          this.currentPartitionIterator = getPartitionIterator((PartitionableDataset) dataset);
        } else {
          WorkUnit workUnit = workUnitForDataset(dataset);
          addDatasetInfoToWorkUnit(workUnit, dataset, this.generatedWorkUnits++);
          return workUnit;
        }
      }

      WorkUnit workUnit = NoopTask.noopWorkunit();
      workUnit.setProp(WORK_UNIT_ORDINAL, this.generatedWorkUnits);

      this.generatedWorkUnits = Integer.MAX_VALUE;

      workUnit.setProp(END_OF_DATASETS_KEY, true);
      return workUnit;
    }

    private void addDatasetInfoToWorkUnit(WorkUnit workUnit, Dataset dataset, int workUnitOrdinal) {
      workUnit.setProp(DATASET_URN, dataset.getUrn());
      workUnit.setProp(WORK_UNIT_ORDINAL, workUnitOrdinal);
    }

    private void addPartitionInfoToWorkUnit(WorkUnit workUnit, PartitionableDataset.DatasetPartition partition) {
      workUnit.setProp(PARTITION_URN, partition.getUrn());
    }
  }

  /**
   * Sort input stream lexicographically. Noop if the input stream is already sorted.
   */
  private <T extends URNIdentified> Stream<T> sortStreamLexicographically(Stream<T> inputStream) {
    Spliterator<T> spliterator = inputStream.spliterator();
    if (spliterator.hasCharacteristics(Spliterator.SORTED) &&
        spliterator.getComparator().equals(this.lexicographicalComparator)) {
      return StreamSupport.stream(spliterator, false);
    }
    return StreamSupport.stream(spliterator, false).sorted(this.lexicographicalComparator);
  }

}
